package trie;

import java.util.LinkedList;
import java.util.List;

/**
 * @Author: yuisama
 * @Date: 2022/10/9 14:46
 * @Description:字典树实现
 * 每个节点标识isEnd字段
 * 字典树的每个树枝表示一个字符 eg : next['a'] != 0，说明当前树枝存储了字符a
 */
public class TrieSet {
    // 节点定义
    private static class TrieNode {
        // 下一个字符内容，eg : next['a'] != null，就说明下个节点保存的是a字符
        private final TrieNode[] next = new TrieNode[R];
        // 当前节点是否是某个单词结尾
        private boolean isEnd;
    }
    // ----------------------------

    // ASCII 字符字典树
    public static final int R = 256;

    private TrieNode root;

    private int size;

    public int size() {
        return size;
    }


    /**
     * 向当前字典树中添加字符串key
     * @param key
     * @return
     */
    public boolean add(String key) {
        if (key == null) {
            return false;
        }
        // 已经存在该字符串
        if (contains(key)) {
            return false;
        }
        root = add(root,key,0);
        size ++;
        return true;
    }

    /**
     * 返回当前字典树中所有满足模式串的字符串，支持通配符"."
     * @param pattern
     * @return
     */
    public List<String> keysWithPattern(String pattern) {
        List<String> res = new LinkedList<>();
        traverse(root,new StringBuilder(),pattern,0,res);
        return res;
    }

    // 在当前以node为根的字典树中，搜索pattern[i..n]
    private void traverse(TrieNode node, StringBuilder path, String pattern, int i, List<String> res) {
        // base case
        if (node == null) {
            return;
        }
        if (i == pattern.length()) {
            if (node.isEnd) {
                // 满足模式串且为合法字符串，加入结果集
                res.add(path.toString());
            }
            return;
        }
        char c = pattern.charAt(i);
        if (c == '.') {
            // 碰到通配符，扫描所有子树
            for (char j = 0; j < R; j++) {
                // 回溯前选择
                path.append(j);
                traverse(node.next[j],path,pattern,i + 1,res);
                // 回溯后撤销选择
                path.deleteCharAt(path.length() - 1);
            }
        }else {
            path.append(c);
            traverse(node.next[c],path,pattern,i + 1,res);
            path.deleteCharAt(path.length() - 1);
        }
    }


    /**
     * 判断当前字典树中是否包含指定prefix的前缀字符串
     * @param prefix
     * @return
     */
    public boolean startsWith(String prefix) {
        TrieNode p = getNode(root,prefix);
        return p != null;
    }

    /**
     * 在当前字典树中是否包含匹配模式串的字符串
     * @param pattern
     * @return
     */
    public boolean matches(String pattern) {
        if (pattern == null) {
            return false;
        }
        return matches(root,pattern,0);
    }

    private boolean matches(TrieNode node, String pattern, int i) {
        // base case
        if (node == null) {
            return false;
        }
        if (i == pattern.length()) {
            return node.isEnd;
        }
        char c = pattern.charAt(i);
        if (c == '.') {
            // 遇到通配符,扫描所有的子树，查找是否有满足模式的子串
            for (char j = 0; j < R; j++) {
                if (matches(node.next[j],pattern,i + 1)) {
                    return true;
                }
            }
        }
        return matches(node.next[c],pattern,i + 1);
    }

    /**
     * 在当前字典树中搜索所有前缀为prefix的字符串并返回
     * @param prefix
     * @return
     */
    public List<String> keysWithPrefix(String prefix) {
        List<String> res = new LinkedList<>();
        if (prefix == null) {
            return res;
        }
        TrieNode node = getNode(root,prefix);
        StringBuilder path = new StringBuilder(prefix);
        // 从node开始向下遍历搜索Trie
        // 多叉树的遍历
        traverse(node,path,res);
        return res;
    }

    /**
     * 多叉树遍历
     * @param node
     * @param path
     * @param res
     */
    private void traverse(TrieNode node, StringBuilder path, List<String> res) {
        if (node == null) {
            return;
        }
        // base case,走到终止节点，添加结果集
        if (node.isEnd) {
            res.add(path.toString());
        }
        // 遍历当前节点的next域，将所有非空的字符添加到路径中
        // 典型的回溯框架，在递归前进行选择，在递归后撤销选择
        for (char i = 0; i < R; i++) {
            // 选择一个字符
            path.append(i);
            // 继续向下搜索
            traverse(node.next[i],path,res);
            // 撤销当前选择
            path.deleteCharAt(path.length() - 1);
        }
    }

    /**
     * 在当前Trie中查找指定字符串query的最短前缀子串
     * @param query
     * @return
     */
    public String shortestPrefixOf(String query) {
        TrieNode p = root;
        int i = 0;
        // 在当前trie中待查找字符串的字符在树中还有字符没有走完
        while (i < query.length() && p != null) {
            char c = query.charAt(i);
            if (p.isEnd) {
                // 走到终止节点，当前路径上的字符就是最短子串
                break;
            }
            p = p.next[c];
            i ++;
        }
        // 只有当p不为空且为终止节点才说明存在满足题意的最短子串
        if (p != null && p.isEnd) {
            // 左闭右开区间
            return query.substring(0,i);
        }
        return "";
    }

    /**
     * 在当前Trie中查找字符串query的最长前缀子串
     * @param query
     * @return
     */
    public String longestPrefixOf(String query) {
        TrieNode p = root;
        // len只有碰到p为终止节点时才更新
        int i = 0,len = 0;
        while (i < query.length() && p != null) {
            if (p.isEnd) {
                len = i;
            }
            char c = query.charAt(i);
            p = p.next[c];
            i ++;
        }
        if (p != null && p.isEnd) {
            // 此时再更新下len，防止query串本身就是最长子串的情况
            len = i;
        }
        return query.substring(0,len);
    }

    /**
     * 在Trie中删除指定字符串
     * @param key
     * @return
     */
    public boolean remove(String key) {
        if (key == null) {
            return false;
        }
        if (!contains(key)) {
            return false;
        }
        root = remove(root,key,0);
        size --;
        return true;
    }

    /**
     * 在Trie中从node节点开始删除字符串key[i..legnth - 1]
     * @param node
     * @param key
     * @param i
     * @return
     */
    private TrieNode remove(TrieNode node, String key, int i) {
        // base case，走到叶子节点，从叶子节点开始回溯删除节点
        if (i == key.length()) {
            node.isEnd = false;
        }else {
            char c = key.charAt(i);
            // 先递归到叶子节点
            node.next[c] = remove(node.next[c],key,i + 1);
        }
        if (node.isEnd) return node;
        for (int j = 0; j < R; j++) {
            if (node.next[j] != null) {
                // 只要有一个非空的节点，都不能删除
                return node;
            }
        }
        // 只有当前节点不是终止节点且next数组全为空时才可删除当前节点
        return null;
    }

    /**
     * 从node节点开始存储字符串key[i...length - 1]，返回插入后的树根
     * @param node
     * @param key
     * @param i
     * @return
     */
    private TrieNode add(TrieNode node, String key, int i) {
        // base case 1.. node为空，创建新节点
        if (node == null) {
            node = new TrieNode();
        }
        // base case 2 字符串已经走到底，将当前node置为isEnd节点并返回
        if (i == key.length()) {
            node.isEnd = true;
            return node;
        }
        // 当前方法要存储的字符
        char c = key.charAt(i);
        // 继续构造剩下的字符，更新当前树根
        node.next[c] = add(node.next[c],key,i + 1);
        return node;
    }

    /**
     * 判断当前字典树中是否包含指定字符串
     * @param key
     * @return
     */
    public boolean contains(String key) {
        if (key == null || key.length() == 0) {
            return false;
        }
        TrieNode x = getNode(root,key);
        // x不为空且为isEnd字段说明包含此字符串
        return x != null && x.isEnd;
    }

    /**
     * 从字典树的node节点开始搜索字符串key，返回搜索后的叶子节点
     * @param node
     * @param key
     * @return
     */
    private TrieNode getNode(TrieNode node, String key) {
        TrieNode p = node;
        for (int i = 0; i < key.length(); i++) {
            char c = key.charAt(i);
            if (p == null) {
                return null;
            }
            p = p.next[c];
        }
        return p;
    }

}
